import torch

import runner1
import dataset
import model1
import assistant

if __name__ == "__main__":
    model = model1.LinearModel(output_num=5)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
    loss_fun = torch.nn.CrossEntropyLoss()
    lr_schedule = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
    train_iter, dev_iter, test_iter = dataset.get_dataloader('../my_spider_douban_formal/output/clean_data.jsonl', 256)
    runner_man = runner1.Runner(model, optimizer, loss_fun, train_iter, logger=assistant.logger, multi=False,
                                device=torch.device('cuda:0'), scheduler=lr_schedule, valid_loader=dev_iter)
    runner_man.run(1000)
